import datetime

import torch
import torch.nn as nn
import torch.nn.functional as F
import os
# from matplotlib import pyplot as plt

class Net(nn.Module):
    def __init__(self, dropout=0.1):
        super(Net, self).__init__()
        self.dropout = nn.Dropout(dropout)
        self.conv1 = nn.Conv2d(1, 10, 5)
        self.conv2 = nn.Conv2d(10, 25, 5)
        self.fc1 = nn.Linear(1 * 25 * 7 * 7, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 34)

    def forward(self, x):
        # 池化出来大小直接除2
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = self.dropout(x)
        x = F.max_pool2d(F.relu(self.conv2(x)), (2, 2))
        x = self.dropout(x)
        x = x.view(-1, self.num_flat_features(x))
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

    def num_flat_features(self, x):
        size = x.size()[1:]  # all dimensions except the batch dimension
        num_features = 1
        for s in size:
            num_features *= s
        return num_features


# net = Net()
# print(net)

import torch.optim as optim
import net_data as data
import torch

# def train():
#     net = Net()
#     opt = optim.SGD(net.parameters(), lr=0.01)
#     epoch = 500
#     batch_size = 50
#     trainloader = data.trainloader(batch_size)
#     for e in range(epoch):
#         for step, d in enumerate(trainloader):
#             opt.zero_grad()
#             out = net(d["data"])
#             lf = nn.CrossEntropyLoss()
#             loss = lf(out, d["label"])
#             loss.backward()
#             opt.step()
#             if (e % 50 == 0):
#                 print("e : {} , step : {}, loss : {}".format(e, step, loss))
#     torch.save(net.state_dict(),"./model/net_gpu.pt")
#     torch.save(opt.state_dict(), "./model/opt.pt")

def train_gpu():
    use_cuda = torch.cuda.is_available()
    if(use_cuda):
        print("use gpu cuda")
    else:
        print("use cpu")

    device = torch.device("cuda:0" if use_cuda else "cpu")
    net = Net()
    net.to(device)

    opt = optim.SGD(net.parameters(), lr=0.01)
    epoch = 200
    batch_size = 50
    epoch_list = []
    loss_list = []
    trainloader = data.trainloader(batch_size)
    st = datetime.datetime.now()
    loss = 0

    for e in range(epoch):
        for step, d in enumerate(trainloader):
            data_cuda =  d["data"].to(device)
            label_cuda = d["label"].to(device)
            opt.zero_grad()
            out = net(data_cuda)
            lf = nn.CrossEntropyLoss()
            loss = lf(out, label_cuda)
            loss.backward()
            opt.step()

            if (e % 50 == 0 or step == 1):
                epoch_list.append(float("{}".format(e)))
                loss_list.append(float("{}".format(loss)))
                print("e : {} , step : {}, loss : {}".format(e, step, loss))

    print("loss : {}".format(loss))
    print("cost time:",datetime.datetime.now() - st)
    # drawLoss(epoch_list,loss_list)
    saveModel(net, opt)

# def drawLoss(epoch_list,loss_list):
#     plt.plot(epoch_list, loss_list)
#     plt.ylabel('loss')
#     plt.xlabel('epoch')
#     plt.show()
#     pass

def saveModel(net,opt):
    net_model_file_path = "model/net.pt"
    opt_net_file_path = "./model/opt.pt"

    if os.path.exists(net_model_file_path):
        bak_path = "./model/net_bak.pt"
        if os.path.exists(bak_path):
            os.remove(bak_path)
        os.rename(net_model_file_path, bak_path)

    if os.path.exists(opt_net_file_path):
        bak_path = "./model/opt_bak.pt"
        if os.path.exists(bak_path):
            os.remove(bak_path)
        os.rename(opt_net_file_path, bak_path)

    torch.save(net.state_dict(), net_model_file_path)
    torch.save(opt.state_dict(), opt_net_file_path)
    pass

if (__name__ == "__main__"):
    # print(torch.__version__)
    # print(torch.cuda.is_available())
    train_gpu()
